// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <boost/random/uniform_int_distribution.hpp>
#include <gtest/gtest.h>

#include <string>

#include <melon/experimental/EventCount.h>

#include <pollux/common/file/file_systems.h>
#include <pollux/common/file/utils.h>
#include <pollux/testing/common/faulty_file_system.h>
#include <pollux/common/hyperloglog/sparse_hll.h>
#include <pollux/common/testutil/test_value.h>
#include <pollux/dwio/dwrf/writer/writer.h>
#include <pollux/exec/operator_trace_reader.h>
#include <pollux/exec/partition_function.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/exec/table_writer.h>
#include <pollux/exec/trace_util.h>
#include <pollux/testing/exec/util/arbitrator_test_util.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/serializers/presto_serializer.h>
#include <pollux/tool/trace/hash_join_replayer.h>
#include <pollux/tool/trace/trace_replay_runner.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::core;
using namespace kumo::pollux::common;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::connector;
using namespace kumo::pollux::connector::hive;
using namespace kumo::pollux::dwio::common;
using namespace kumo::pollux::common::testutil;
using namespace kumo::pollux::common::hll;
using namespace kumo::pollux::tests::utils;

namespace kumo::pollux::tool::trace::test {
class HashJoinReplayerTest : public HiveConnectorTestBase {
 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
    HiveConnectorTestBase::SetUpTestCase();
    registerFaultyFileSystem();
    if (!isRegisteredVectorSerde()) {
      serializer::presto::PrestoVectorSerde::registerVectorSerde();
    }
    Type::registerSerDe();
    common::Filter::registerSerDe();
    connector::hive::HiveTableHandle::registerSerDe();
    connector::hive::LocationHandle::registerSerDe();
    connector::hive::HiveColumnHandle::registerSerDe();
    connector::hive::HiveInsertTableHandle::registerSerDe();
    connector::hive::HiveInsertFileNameGenerator::registerSerDe();
    connector::hive::HiveConnectorSplit::registerSerDe();
    core::PlanNode::registerSerDe();
    core::ITypedExpr::registerSerDe();
    registerPartitionFunctionSerDe();
  }

  void TearDown() override {
    probeInput_.clear();
    buildInput_.clear();
    HiveConnectorTestBase::TearDown();
  }

  struct PlanWithSplits {
    core::PlanNodePtr plan;
    core::PlanNodeId probeScanId;
    core::PlanNodeId buildScanId;
    melon::F14FastMap<core::PlanNodeId, std::vector<exec::Split>> splits;

    explicit PlanWithSplits(
        const core::PlanNodePtr& _plan,
        const core::PlanNodeId& _probeScanId = "",
        const core::PlanNodeId& _buildScanId = "",
        const melon::F14FastMap<
            core::PlanNodeId,
            std::vector<pollux::exec::Split>>& _splits = {})
        : plan(_plan),
          probeScanId(_probeScanId),
          buildScanId(_buildScanId),
          splits(_splits) {}
  };

  RowTypePtr concat(const RowTypePtr& a, const RowTypePtr& b) {
    std::vector<std::string> names = a->names();
    std::vector<TypePtr> types = a->children();

    for (auto i = 0; i < b->size(); ++i) {
      names.push_back(b->nameOf(i));
      types.push_back(b->childAt(i));
    }

    return ROW(std::move(names), std::move(types));
  }

  std::vector<RowVectorPtr>
  makeVectors(int32_t count, int32_t rowsPerVector, const RowTypePtr& rowType) {
    return HiveConnectorTestBase::makeVectors(rowType, count, rowsPerVector);
  }

  std::vector<Split> makeSplits(
      const std::vector<RowVectorPtr>& inputs,
      const std::string& path,
      memory::MemoryPool* writerPool) {
    std::vector<Split> splits;
    for (auto i = 0; i < 4; ++i) {
      const std::string filePath = fmt::format("{}/{}", path, i);
      writeToFile(filePath, inputs);
      splits.emplace_back(makeHiveConnectorSplit(filePath));
    }

    return splits;
  }

  PlanWithSplits createPlan(
      const std::string& tableDir,
      core::JoinType joinType,
      const std::vector<std::string>& probeKeys,
      const std::vector<std::string>& buildKeys,
      const std::vector<RowVectorPtr>& probeInput,
      const std::vector<RowVectorPtr>& buildInput) {
    auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
    const std::vector<Split> probeSplits =
        makeSplits(probeInput, fmt::format("{}/probe", tableDir), pool());
    const std::vector<Split> buildSplits =
        makeSplits(buildInput, fmt::format("{}/build", tableDir), pool());
    core::PlanNodeId probeScanId;
    core::PlanNodeId buildScanId;
    const auto outputColumns = concat(
                                   as_row_type(probeInput_[0]->type()),
                                   as_row_type(buildInput_[0]->type()))
                                   ->names();
    auto plan = PlanBuilder(planNodeIdGenerator)
                    .tableScan(probeType_)
                    .capturePlanNodeId(probeScanId)
                    .hashJoin(
                        probeKeys,
                        buildKeys,
                        PlanBuilder(planNodeIdGenerator)
                            .tableScan(buildType_)
                            .capturePlanNodeId(buildScanId)
                            .planNode(),
                        /*filter=*/"",
                        outputColumns,
                        joinType,
                        false)
                    .capturePlanNodeId(traceNodeId_)
                    .planNode();
    return PlanWithSplits{
        plan,
        probeScanId,
        buildScanId,
        {{probeScanId, probeSplits}, {buildScanId, buildSplits}}};
  }

  core::PlanNodeId traceNodeId_;
  RowTypePtr probeType_{
      ROW({"t0", "t1", "t2", "t3"}, {BIGINT(), VARCHAR(), SMALLINT(), REAL()})};

  RowTypePtr buildType_{
      ROW({"u0", "u1", "u2", "u3"},
          {BIGINT(), INTEGER(), SMALLINT(), VARCHAR()})};
  std::vector<RowVectorPtr> probeInput_ = makeVectors(5, 100, probeType_);
  std::vector<RowVectorPtr> buildInput_ = makeVectors(3, 100, buildType_);

  const std::vector<std::string> probeKeys_{"t0"};
  const std::vector<std::string> buildKeys_{"u0"};
  const std::shared_ptr<TempDirectoryPath> testDir_ =
      TempDirectoryPath::create();
  const std::string tableDir_ =
      fmt::format("{}/{}", testDir_->getPath(), "table");
};

TEST_F(HashJoinReplayerTest, basic) {
  const auto planWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder builder(planWithSplits.plan);
  for (const auto& [planNodeId, nodeSplits] : planWithSplits.splits) {
    builder.splits(planNodeId, nodeSplits);
  }
  const auto result = builder.copyResults(pool());

  const auto traceRoot =
      fmt::format("{}/{}/traceRoot/", testDir_->getPath(), "basic");
  std::shared_ptr<Task> task;
  auto tracePlanWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder traceBuilder(tracePlanWithSplits.plan);
  traceBuilder.maxDrivers(4)
      .config(core::QueryConfig::kQueryTraceEnabled, true)
      .config(core::QueryConfig::kQueryTraceDir, traceRoot)
      .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
      .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
      .config(core::QueryConfig::kQueryTraceNodeIds, traceNodeId_);
  for (const auto& [planNodeId, nodeSplits] : tracePlanWithSplits.splits) {
    traceBuilder.splits(planNodeId, nodeSplits);
  }
  auto traceResult = traceBuilder.copyResults(pool(), task);

  assertEqualResults({result}, {traceResult});

  const auto taskId = task->taskId();
  const auto replayingResult = HashJoinReplayer(
                                   traceRoot,
                                   task->queryCtx()->queryId(),
                                   task->taskId(),
                                   traceNodeId_,
                                   "HashJoin",
                                   "",
                                   0,
                                   executor_.get())
                                   .run();
  assertEqualResults({result}, {replayingResult});
}

TEST_F(HashJoinReplayerTest, partialDriverIds) {
  const std::shared_ptr<TempDirectoryPath> testDir =
      TempDirectoryPath::create(true);
  const std::string tableDir =
      fmt::format("{}/{}", testDir->getPath(), "table");
  const auto planWithSplits = createPlan(
      tableDir,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder builder(planWithSplits.plan);
  for (const auto& [planNodeId, nodeSplits] : planWithSplits.splits) {
    builder.splits(planNodeId, nodeSplits);
  }
  const auto result = builder.copyResults(pool());

  const auto traceRoot =
      fmt::format("{}/{}/traceRoot/", testDir->getPath(), "basic");
  std::shared_ptr<Task> task;
  auto tracePlanWithSplits = createPlan(
      tableDir,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder traceBuilder(tracePlanWithSplits.plan);
  traceBuilder.maxDrivers(4)
      .config(core::QueryConfig::kQueryTraceEnabled, true)
      .config(core::QueryConfig::kQueryTraceDir, traceRoot)
      .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
      .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
      .config(core::QueryConfig::kQueryTraceNodeIds, traceNodeId_);
  for (const auto& [planNodeId, nodeSplits] : tracePlanWithSplits.splits) {
    traceBuilder.splits(planNodeId, nodeSplits);
  }
  auto traceResult = traceBuilder.copyResults(pool(), task);

  assertEqualResults({result}, {traceResult});

  const auto taskId = task->taskId();
  const auto taskTraceDir =
      exec::trace::getTaskTraceDirectory(traceRoot, *task);
  const auto opTraceDir =
      exec::trace::getOpTraceDirectory(taskTraceDir, traceNodeId_, 0, 0);
  const auto opTraceDataFile = exec::trace::getOpTraceInputFilePath(opTraceDir);
  auto faultyFs = faultyFileSystem();
  faultyFs->setFileInjectionHook([&](FaultFileOperation* op) {
    if ((op->type == FaultFileOperation::Type::kRead ||
         op->type == FaultFileOperation::Type::kReadv) &&
        op->path == opTraceDataFile) {
      POLLUX_FAIL("Read wrong data file {}", opTraceDataFile);
    }
  });

  if (faultyFs->openFileForRead(opTraceDataFile)->size() > 0) {
    POLLUX_ASSERT_THROW(
        HashJoinReplayer(
            traceRoot,
            task->queryCtx()->queryId(),
            task->taskId(),
            traceNodeId_,
            "HashJoin",
            "0",
            0,
            executor_.get())
            .run(),
        "Read wrong data file");
  }
  HashJoinReplayer(
      traceRoot,
      task->queryCtx()->queryId(),
      task->taskId(),
      traceNodeId_,
      "HashJoin",
      "1,3",
      0,
      executor_.get())
      .run();
  faultyFs->clearFileFaultInjections();
}

TEST_F(HashJoinReplayerTest, runner) {
  const auto testDir = TempDirectoryPath::create();
  const auto traceRoot = fmt::format("{}/{}", testDir->getPath(), "traceRoot");
  std::shared_ptr<Task> task;
  auto tracePlanWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder traceBuilder(tracePlanWithSplits.plan);
  traceBuilder.config(core::QueryConfig::kQueryTraceEnabled, true)
      .config(core::QueryConfig::kQueryTraceDir, traceRoot)
      .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
      .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
      .config(core::QueryConfig::kQueryTraceNodeIds, traceNodeId_);
  for (const auto& [planNodeId, nodeSplits] : tracePlanWithSplits.splits) {
    traceBuilder.splits(planNodeId, nodeSplits);
  }
  auto traceResult = traceBuilder.copyResults(pool(), task);

  const auto taskTraceDir =
      exec::trace::getTaskTraceDirectory(traceRoot, *task);
  const auto probeOperatorTraceDir = exec::trace::getOpTraceDirectory(
      taskTraceDir,
      traceNodeId_,
      /*pipelineId=*/0,
      /*driverId=*/0);
  const auto probeSummary =
      exec::trace::OperatorTraceSummaryReader(probeOperatorTraceDir, pool())
          .read();
  ASSERT_EQ(probeSummary.opType, "HashProbe");
  ASSERT_GT(probeSummary.peakMemory, 0);
  ASSERT_GT(probeSummary.inputRows, 0);
  ASSERT_GT(probeSummary.inputBytes, 0);
  ASSERT_EQ(probeSummary.rawInputRows, 0);
  ASSERT_EQ(probeSummary.rawInputBytes, 0);

  const auto buildOperatorTraceDir = exec::trace::getOpTraceDirectory(
      taskTraceDir,
      traceNodeId_,
      /*pipelineId=*/1,
      /*driverId=*/0);
  const auto buildSummary =
      exec::trace::OperatorTraceSummaryReader(buildOperatorTraceDir, pool())
          .read();
  ASSERT_EQ(buildSummary.opType, "HashBuild");
  ASSERT_GT(buildSummary.peakMemory, 0);
  ASSERT_GT(buildSummary.inputRows, 0);
  // NOTE: the input bytes is 0 because of the lazy materialization.
  ASSERT_EQ(buildSummary.inputBytes, 0);
  ASSERT_EQ(buildSummary.rawInputRows, 0);
  ASSERT_EQ(buildSummary.rawInputBytes, 0);

  FLAGS_root_dir = traceRoot;
  FLAGS_query_id = task->queryCtx()->queryId();
  FLAGS_task_id = task->taskId();
  FLAGS_node_id = traceNodeId_;
  FLAGS_summary = true;
  {
    TraceReplayRunner runner;
    runner.init();
    runner.run();
  }

  FLAGS_task_id = task->taskId();
  FLAGS_driver_ids = "";
  FLAGS_summary = false;
  {
    TraceReplayRunner runner;
    runner.init();
    runner.run();
  }
}

DEBUG_ONLY_TEST_F(HashJoinReplayerTest, hashBuildSpill) {
  const auto planWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);
  AssertQueryBuilder builder(planWithSplits.plan);
  for (const auto& [planNodeId, nodeSplits] : planWithSplits.splits) {
    builder.splits(planNodeId, nodeSplits);
  }
  const auto result = builder.copyResults(pool());

  const auto traceRoot =
      fmt::format("{}/{}/traceRoot/", testDir_->getPath(), "hash_build_spill");
  const auto spillDir =
      fmt::format("{}/{}/spillDir/", testDir_->getPath(), "hash_build_spill");
  std::shared_ptr<Task> task;
  auto tracePlanWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);

  std::atomic_bool injectSpillOnce{true};
  SCOPED_TESTVALUE_SET(
      "kumo::pollux::exec::HashBuild::finishHashBuild",
      std::function<void(Operator*)>([&](Operator* op) {
        if (!injectSpillOnce.exchange(false)) {
          return;
        }
        Operator::ReclaimableSectionGuard guard(op);
        testingRunArbitration(op->pool());
      }));

  AssertQueryBuilder traceBuilder(tracePlanWithSplits.plan);
  traceBuilder.config(core::QueryConfig::kQueryTraceEnabled, true)
      .config(core::QueryConfig::kQueryTraceDir, traceRoot)
      .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
      .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
      .config(core::QueryConfig::kQueryTraceNodeIds, traceNodeId_)
      .config(core::QueryConfig::kSpillEnabled, true)
      .config(core::QueryConfig::kJoinSpillEnabled, true)
      .spillDirectory(spillDir);
  for (const auto& [planNodeId, nodeSplits] : tracePlanWithSplits.splits) {
    traceBuilder.splits(planNodeId, nodeSplits);
  }

  auto traceResult = traceBuilder.copyResults(pool(), task);
  auto taskStats = exec::toPlanStats(task->taskStats());
  const auto& stats = taskStats.at(traceNodeId_);
  auto opStats = toOperatorStats(task->taskStats());
  ASSERT_GT(
      opStats.at("HashBuild").runtimeStats[Operator::kSpillWrites].sum, 0);
  ASSERT_GT(stats.spilledBytes, 0);
  ASSERT_GT(stats.spilledRows, 0);
  ASSERT_GT(stats.spilledFiles, 0);
  ASSERT_GT(stats.spilledPartitions, 0);

  assertEqualResults({result}, {traceResult});

  const auto taskId = task->taskId();
  const auto replayingResult = HashJoinReplayer(
                                   traceRoot,
                                   task->queryCtx()->queryId(),
                                   task->taskId(),
                                   traceNodeId_,
                                   "HashJoin",
                                   "",
                                   0,
                                   executor_.get())
                                   .run();
  assertEqualResults({result}, {replayingResult});
}

DEBUG_ONLY_TEST_F(HashJoinReplayerTest, hashProbeSpill) {
  const auto planWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);

  AssertQueryBuilder builder(planWithSplits.plan);
  for (const auto& [planNodeId, nodeSplits] : planWithSplits.splits) {
    builder.splits(planNodeId, nodeSplits);
  }
  const auto result = builder.copyResults(pool());

  const auto traceRoot =
      fmt::format("{}/{}/traceRoot/", testDir_->getPath(), "hash_probe_spill");
  const auto spillDir =
      fmt::format("{}/{}/spillDir/", testDir_->getPath(), "hash_probe_spill");
  std::shared_ptr<Task> task;
  auto tracePlanWithSplits = createPlan(
      tableDir_,
      core::JoinType::kInner,
      probeKeys_,
      buildKeys_,
      probeInput_,
      buildInput_);

  std::atomic_bool injectProbeSpillOnce{true};
  SCOPED_TESTVALUE_SET(
      "kumo::pollux::exec::Driver::runInternal::getOutput",
      std::function<void(Operator*)>([&](Operator* op) {
        if (!isHashProbeMemoryPool(*op->pool())) {
          return;
        }
        if (!injectProbeSpillOnce.exchange(false)) {
          return;
        }
        testingRunArbitration(op->pool());
      }));

  AssertQueryBuilder traceBuilder(tracePlanWithSplits.plan);
  traceBuilder.config(core::QueryConfig::kQueryTraceEnabled, true)
      .config(core::QueryConfig::kQueryTraceDir, traceRoot)
      .config(core::QueryConfig::kQueryTraceMaxBytes, 100UL << 30)
      .config(core::QueryConfig::kQueryTraceTaskRegExp, ".*")
      .config(core::QueryConfig::kQueryTraceNodeIds, traceNodeId_)
      .config(core::QueryConfig::kSpillEnabled, true)
      .config(core::QueryConfig::kJoinSpillEnabled, true)
      .spillDirectory(spillDir);
  for (const auto& [planNodeId, nodeSplits] : tracePlanWithSplits.splits) {
    traceBuilder.splits(planNodeId, nodeSplits);
  }

  auto traceResult = traceBuilder.copyResults(pool(), task);
  auto taskStats = exec::toPlanStats(task->taskStats());
  const auto& stats = taskStats.at(traceNodeId_);
  auto opStats = toOperatorStats(task->taskStats());
  ASSERT_GT(
      opStats.at("HashProbe").runtimeStats[Operator::kSpillWrites].sum, 0);

  ASSERT_GT(stats.spilledBytes, 0);
  ASSERT_GT(stats.spilledRows, 0);
  ASSERT_GT(stats.spilledFiles, 0);
  ASSERT_GT(stats.spilledPartitions, 0);

  assertEqualResults({result}, {traceResult});

  const auto taskId = task->taskId();
  const auto replayingResult = HashJoinReplayer(
                                   traceRoot,
                                   task->queryCtx()->queryId(),
                                   task->taskId(),
                                   traceNodeId_,
                                   "HashJoin",
                                   "",
                                   0,
                                   executor_.get())
                                   .run();
  assertEqualResults({result}, {replayingResult});
}
} // namespace kumo::pollux::tool::trace::test
